from langchain_core.runnables import RunnableConfig

from owl_ai.graph.graph_node import BaseGraphNode


class GraphConditionRouter(BaseGraphNode):
    """
    图条件路由
    """

    def __init__(self, node_config: dict):
        self.conditions = node_config.get("conditions")
        self.input = node_config.get("input")

    def __call__(self, state, config: RunnableConfig):
        input_params = self.input_params_assem(state, self.input)
        for condition in self.conditions:
            condition_type = condition.get("conditionType")
            condition_value = condition.get("conditionValue")
            if condition_type == "eq":
                if input_params.get("conditionValue") == condition_value:
                    return condition.get("nextNode")
        return None
